from typing import Optional, Sequence, Tuple

import numpy as np
import util
from timer import timers

def training_dp(num_modules,
                num_devices,
                num_microbatches,
                submesh_choices,
                compute_cost,
                max_n_succ_stages):
    """Auto stage dynamic programming."""
    timers("stage-construction-dp").start()

    all_possible_stage_costs = np.sort(np.unique(compute_cost))
    best_cost = np.inf  # 初始化优化目标为无穷大
    best_solution = None
    last_max_stage_cost = 0.0  # 记录上一次使用的最大阶段成本

    gap = 1e-6
    # 检查 all_possible_stage_costs 非空
    assert len(all_possible_stage_costs), "no solution in auto stage construction."

    # 论文中枚举并固定 t_max 的操作
    for max_stage_cost in all_possible_stage_costs:
        # 论文中 early pruning 的两个操作
        if max_stage_cost * num_microbatches >= best_cost:
            break
        if max_stage_cost - last_max_stage_cost < gap:
            continue
        cost, solution = training_dp_impl(num_modules,
                                          num_devices,
                                          num_microbatches,
                                          submesh_choices,
                                          compute_cost,
                                          max_n_succ_stages,
                                          max_stage_cost)
        if cost < best_cost:
            best_cost = cost
            best_solution = solution
        last_max_stage_cost = max_stage_cost

    print("------------wjq---------")
    print("best solution is : ", best_solution) # [((0, 5), 1, 2), ((5, 8), 3, 4), ((8, 9), 2, 4), ((9, 16), 1, 0)]
    # submesh_choice
    # autosharding_choice
    print("------------wjq---------")

    timers("stage-construction-dp").stop()
    return best_cost, best_solution

@util.maybe_numba_jit
def training_dp_impl(num_modules,
                     num_devices,
                     num_microbatches,
                     submesh_choices,
                     compute_cost,
                     max_n_succ_stages,
                     max_stage_cost):
    """The core implementation of the DP algorithm."""
    # 初始化为无穷大
    f = np.full((num_modules + 1, num_modules + 1, num_devices + 1),
                np.inf,
                dtype=np.float32)
    # 每个stage对应的t_max
    f_stage_max = np.full((num_modules + 1, num_modules + 1, num_devices + 1),
                          0.0,
                          dtype=np.float32)
    # f_argmin 用于存储使得 f 取得最小值的索引，用于回溯最优路径。
    # 它的每个元素是一个三元组 (k, m, n_config)，表示下一个开始层、子网格选择和自动分片配置
    f_argmin = np.full(
        (num_modules + 1,
         num_modules + 1,
         num_devices + 1,
         3),
        -1,
        dtype=np.int32)

    # 初始化dp数组
    # f(s, k, d; t_max)
    f[0, num_modules, 0] = 0

    # 循环遍历流水线的阶段数量，从 1 到 num_layers。每个阶段都需要尝试不同的划分   s
    for s in range(1, num_modules + 1):
        # 表示从最后一层开始倒序遍历层数，i 表示当前阶段开始的层数  i
        for i in range(num_modules - 1, -1, -1):
            # 表示当前使用的设备数量，从 1 到 num_devices。每次需要考虑在这个设备数量下如何划分阶段  d
            for d in range(1, num_devices + 1):
                # 表示阶段的结束层，从 i 开始的层到模型的末尾层  k
                for k in range(num_modules, i, -1):
                    # 循环遍历子网格配置，计算子网格中设备的数量 n_submesh_devices，并判断它是否小于等于当前可用设备数 j
                    for m, submesh in enumerate(submesh_choices):
                        n_submesh_devices = np.prod(np.array(submesh))
                        if n_submesh_devices <= d:
                            # TODO(zhuohan): This level of for loop is not
                            #   necessary. It can be optimized by sorting
                            #   the logical mesh shapes.

                            # 循环遍历自动分片的配置
                            for n_config in range(1):
                                # 检查当前阶段划分是否符合允许的最大连续阶段数限制，
                                # 即前一个阶段的长度不能超过 max_n_succ_stages 规定的限制
                                if s - 1 <= max_n_succ_stages[i, k - 1, m,
                                n_config]:
                                    # 获取当前层到 k-1 层的计算成本（由 compute_cost 提供，基于子网格选择和分片配置
                                    stage_cost = compute_cost[i, k - 1, m, n_config]
                                    # 计算新的总成本 new_cost：它等于前一个阶段的成本加上当前阶段的计算成本
                                    new_cost = f[s - 1, k, d - n_submesh_devices] + stage_cost
                                    # 条件满足，更新最小总成本
                                    if stage_cost <= max_stage_cost and new_cost < f[s, i, d]:
                                        f[s, i, d] = new_cost
                                        # 更新当前阶段的最大计算成本，以便后续计算总成本
                                        f_stage_max[s, i, d] = max(float(f_stage_max[s - 1, k, d - n_submesh_devices]),
                                                                   stage_cost)
                                        # 记录用于回溯最优路径的索引（包括下一个开始层、子网格选择和分片配置）
                                        f_argmin[s, i, d] = (k, m, n_config)

    # 遍历所有阶段，找到成本最低的 best_s 和最优总成本 best_total_cost
    best_s = -1
    best_total_cost = np.inf
    for s in range(1, num_modules + 1):
        if f[s, 0, num_devices] < best_total_cost:
            best_s = s
            best_total_cost = f[s, 0, num_devices]
    # 如果最优成本仍然是无穷大，说明没有找到可行的方案，返回无穷大和 None
    if np.isinf(best_total_cost):
        return np.inf, None

    # 计算最终总成本：考虑流水线延迟，即最大阶段成本乘以 num_microbatches - 1
    total_cost = f[best_s, 0, num_devices] + (num_microbatches - 1) * f_stage_max[best_s, 0, num_devices]

    current_s = best_s
    current_module = 0
    current_devices = num_devices

    # 通过 f_argmin 中记录的索引值找到下一个阶段的开始层、子网格选择和自动分片配置
    res = []
    while current_s > 0 and current_module < num_modules and current_devices > 0:
        next_start_module, submesh_choice, num_config = (
            f_argmin[current_s, current_module, current_devices])
        assert next_start_module != -1 and current_devices != -1
        res.append(
            ((int(current_module), int(next_start_module)), int(submesh_choice),
             1)
        )
        current_s -= 1
        current_module = next_start_module
        current_devices -= np.prod(np.array(submesh_choices[submesh_choice]))
    assert (current_s == 0 and current_module == num_modules and
            current_devices == 0)

    return total_cost, res

@util.maybe_numba_jit
def training_dp_impl_2(num_layers, num_devices, submesh_sizes,
                       valid_idxs_and_costs, max_n_succ_stages):
    f = np.full((num_layers + 1, num_layers + 1, num_devices + 1),
                np.inf,
                dtype=np.float32)
    f_stage_max = np.full((num_layers + 1, num_layers + 1, num_devices + 1),
                          0.0,
                          dtype=np.float32)
    f_argmin = np.full((num_layers + 1, num_layers + 1, num_devices + 1, 3),
                       -1,
                       dtype=np.int32)
    f[0, num_layers, 0] = 0
    for d in range(1, num_devices + 1):
        for l, i, submesh_id, n_config, stage_cost in valid_idxs_and_costs:
            l, i, submesh_id, n_config = map(int, (l, i, submesh_id, n_config))
            n_submesh_devices = submesh_sizes[submesh_id]
            if n_submesh_devices <= d:
                for s in range(1, num_layers + 1):
                    if s - 1 > max_n_succ_stages[l, i, submesh_id, n_config]:
                        continue

                    new_cost = f[s - 1, i + 1,
                                 d - n_submesh_devices] + stage_cost
                    if new_cost < f[s, l, d]:
                        f[s, l, d] = new_cost
                        f_argmin[s, l, d] = (i + 1, submesh_id, n_config)
                        f_stage_max[s, l, d] = max(float(f_stage_max[s - 1, i + 1, d - n_submesh_devices]), stage_cost)

    return f, f_stage_max, f_argmin


def training_dp_2(
        num_devices,
        num_microbatches,
        submesh_choices,
        compute_cost,
        max_n_succ_stages,
):
    """Faster implementation of the training DP algorihtm."""
    # TODO(zhuohan): Further verify the correctness of this implementation.
    timers("stage-construction-dp").start()

    num_layers = len(compute_cost)
    all_possible_stage_costs = np.sort(np.unique(compute_cost))
    best_cost = np.inf
    best_solution = None
    last_max_stage_cost = 0.0
    # FIXME(zhuohan): Set this gap as a tunable parameter in global config
    gap = 1e-6
    assert len(
        all_possible_stage_costs), "no solution in auto stage construction."

    submesh_sizes = np.array([n * m for (n, m) in submesh_choices],
                             dtype=np.int64)

    for max_stage_cost in all_possible_stage_costs:
        if max_stage_cost - last_max_stage_cost < gap:
            continue
        if max_stage_cost * num_microbatches >= best_cost:
            break

        # Lifts check for stage_cost <= t_max_stage_cost out of the inner dp
        # loop.
        valid_cost_idxs = np.transpose(
            (compute_cost <= max_stage_cost).nonzero())
        # This corresponds to the i of k <= i <= K from eqn. 3 in the alpa
        # paper.
        valid_cost_idxs = valid_cost_idxs[
            valid_cost_idxs[:, 0] <= valid_cost_idxs[:, 1]]
        if len(valid_cost_idxs) == 0:
            continue
        valid_costs = compute_cost[tuple(valid_cost_idxs.T)]
        valid_idxs_and_costs = np.hstack(
            [valid_cost_idxs, valid_costs[:, np.newaxis]])
        # Sort by descending layer idx because DP initializes
        # F[0, num_layers, 0] = 0
        valid_idxs_and_costs = valid_idxs_and_costs[np.flip(
            valid_cost_idxs[:, 1].argsort())]

        # Don't perform backtracking each time (do it only for the best
        # solution).
        f, f_stage_max, f_argmin = training_dp_impl_2(
            num_layers,
            num_devices,
            submesh_sizes,
            valid_idxs_and_costs,
            max_n_succ_stages,
        )

        best_s = f[:, 0, num_devices].argmin()
        best_total_cost = f[best_s, 0, num_devices]
        if np.isinf(best_total_cost):
            continue
        stage_cost = (num_microbatches - 1) * f_stage_max[best_s, 0,
        num_devices]

        if best_total_cost + stage_cost < best_cost:
            best_cost = best_total_cost + stage_cost
            best_solution = best_s, f_argmin
        last_max_stage_cost = max_stage_cost

    assert best_solution is not None, (
        "Unable to find any solution to inter-op dp.")
    best_s, f_argmin = best_solution
    best_solution = get_optimal_submeshes(best_s, f_argmin, num_devices,
                                          num_layers, submesh_sizes)

    print("------------wjq---------")
    print("best solution is : ", best_solution)
    print("best cost is : ", best_cost)
    print("------------wjq---------")

    timers("stage-construction-dp").stop()
    return best_cost, best_solution

def get_submesh_choices(num_hosts: int, num_devices_per_host: int, space: str, manually_specified_submeshes: Optional[Sequence[Tuple[int,int]]] = None):
    """Gets the valid choices of submesh shapes."""
    # if global_config.overwrite_submesh_choices is not None:
    #     return global_config.overwrite_submesh_choices
    submesh_choices = []

    # smaller submeshes:
    i = 1
    while i <= num_devices_per_host:
        submesh_choices.append((1, i))
        i *= 2
    assert submesh_choices[-1][1] == num_devices_per_host, (
        "Only supports the cases where num_devices_per_host is power of two, "
        f"while now num_devices_per_host = {num_devices_per_host}")

    # larger meshes:
    if space == "all":
        for i in range(2, num_hosts + 1):
            submesh_choices.append((i, num_devices_per_host))
    elif space == "power_of_two":
        i = 2
        while i <= num_hosts:
            submesh_choices.append((i, num_devices_per_host))
            i *= 2
    elif space == "small_power_of_two":
        i = 2
        while i <= min(num_hosts, 4):
            submesh_choices.append((i, num_devices_per_host))
            i *= 2
    elif space == "manual":
        submesh_choices = manually_specified_submeshes
    else:
        raise ValueError(f"Invalid submesh space: {space}")

    print(tuple(submesh_choices))
    return tuple(submesh_choices)

@util.maybe_numba_jit
def get_optimal_submeshes(best_s, f_argmin, num_devices, num_layers,
                          submesh_n_devices):
    current_s = best_s
    current_layer = 0
    current_devices = num_devices

    res = []
    while current_s > 0 and current_layer < num_layers and current_devices > 0:
        next_start_layer, submesh_choice, _ = (
            f_argmin[current_s, current_layer, current_devices])
        assert next_start_layer != -1 and current_devices != -1
        print("wjq---子网选择：", submesh_choice)

        res.append(((current_layer, next_start_layer), submesh_choice,
                    1))
        current_s -= 1
        current_layer = next_start_layer
        current_devices -= submesh_n_devices[submesh_choice]
    assert (current_s == 0 and current_layer == num_layers and
            current_devices == 0)

    return res